#!/bin/env  python

import sys

path = "/home/spider/atscript/lib/"
sys.path.append(path)

from atscript import *
from hadoop import *
from assert_func import *

conf = config("rankget.yaml")
sample = base(conf.get("CONF"), conf.get("SCHEDULE"), conf.get("JOBLIST"))
hadoop = hadoop(conf)

def setup():
    if not os.path.exists(conf.get("JOBLIST")+".bak"):
        p = subprocess.Popen("cp %s %s.bak" %(conf.get("JOBLIST"), conf.get("JOBLIST")), shell=True)  
        p.wait()
        cmd = "cat  " + conf.get("JOBLIST") + ".bak" +" | map_only.awk > " + conf.get("JOBLIST")
        os.system(cmd)

def teardown():
    if os.path.exists(conf.get("JOBLIST")+".bak"):
        p = subprocess.Popen("mv %s.bak %s" %(conf.get("JOBLIST"), conf.get("JOBLIST")), shell=True)  
        p.wait()


class Test_rg_getsite_map(at_case):
    def __init__(self, casename, base, jobname, input, output):
        at_case.__init__(self, casename, base, jobname, input, output)
    
    def set_input(self):
        line1 = "notmultidomain.1 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line2 = "notmultidomain.2 0 0 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line3 = "notmultidomain.3 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 2"
        
        line4 = ".multidomain.1 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line5 = ".multidomain.2 0 0 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line6 = ".multidomain.3 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 2"
        
        input_content = "\n".join([line1, line2, line3, line4, line5, line6])
        
        temp_input_file = self.casename + ".input"
        try:
            file_handle = open(temp_input_file, 'w')
            file_handle.write(input_content)
            file_handle.write("\n")
        except IOError, e:
            raise IOError
        finally:
            file_handle.close()

        hadoop.rmfile("/user/spider/at-script/rankget/input/" + temp_input_file)
        hadoop.put(temp_input_file, "/user/spider/at-script/rankget/input/")

        file_handle = open("select_name_only","w")
        file_handle.write("/user/spider/at-script/rankget/input/" + temp_input_file + "\n")
        file_handle.close()

        hadoop.rmfile(self.input + "select_name_only")
        hadoop.put("select_name_only", self.input)
    
    def run(self):
        hadoop.rmdir(self.output)
        hadoop.mkdir(self.output)
        
        print self.run_cmd
        self.execute_cmd(self.run_cmd, True)

    def check_result(self):
        expect_out = self.casename + ".expect_out"
        real_out = self.casename + ".real_out"

        self.execute_cmd("rm -f %s" % real_out)        
        hadoop.getmerge(self.output, real_out)
          
        try:
            file_handle = open(expect_out, 'w')
            file_handle.write("notmultidomain.1\t99999\tNOTMULTIDOMAIN\n")
            file_handle.write("notmultidomain.2\t99999\tNOTMULTIDOMAIN\n")
            file_handle.write("notmultidomain.3\t100001\tNOTMULTIDOMAIN\n")
            file_handle.write(".multidomain.1\t99999\tMULTIDOMAIN\n")
            file_handle.write(".multidomain.2\t99999\tMULTIDOMAIN\n")
            file_handle.write(".multidomain.3\t100001\tMULTIDOMAIN\n")
        except IOError, e:
            raise IOError
        finally:
            file_handle.close()
        
        try:
           assert_file_equal(expect_out, real_out)        
        except AssertionError, e:
           print >>sys.stdout, "case %s failed" % self.casename
           print >>sys.stdout, "%s not equal to %s" % (real_out, expect_out)
           sys.exit(1)
        print >>sys.stdout, "case %s passed" % self.casename


class Test_rg_getsite(at_case):
    def __init__(self, casename, base, jobname, input, output):
        at_case.__init__(self, casename, base, jobname, input, output)
    
    def set_input(self):
        line1 = "notmultidomain.1 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line2 = "notmultidomain.2 0 0 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line3 = "notmultidomain.3 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 2"
        
        line4 = ".multidomain.1 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line5 = ".multidomain.2 0 0 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
        line6 = ".multidomain.3 99999 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 2"
        
        input_content = "\n".join([line1, line2, line3, line4, line5, line6])
        
        temp_input_file = self.casename + ".input"
        try:
            file_handle = open(temp_input_file, 'w')
            file_handle.write(input_content)
            file_handle.write("\n")
        except IOError, e:
            raise IOError
        finally:
            file_handle.close()

        hadoop.rmfile("/user/spider/at-script/rankget/input/" + temp_input_file)
        hadoop.put(temp_input_file, "/user/spider/at-script/rankget/input/")

        file_handle = open("select_name_only","w")
        file_handle.write("/user/spider/at-script/rankget/input/" + temp_input_file + "\n")
        file_handle.close()

        hadoop.rmfile(self.input + "select_name_only")
        hadoop.put("select_name_only", self.input)
   
    def run(self):
        hadoop.rmdir(self.output)
        hadoop.mkdir(self.output)
        
        print self.run_cmd
        self.execute_cmd(self.run_cmd, True)

    def check_result(self):
        expect_out = self.casename + ".expect_out"
        real_out = self.casename + ".real_out"

        self.execute_cmd("rm -f %s" % real_out)        
        hadoop.getmerge(self.output, real_out)
          
        try:
            file_handle = open(expect_out, 'w')
            file_handle.write(".multidomain.3\tMULTIDOMAIN\n")
            file_handle.write("notmultidomain.3\tNOTMULTIDOMAIN\n")
        except IOError, e:
            raise IOError
        finally:
            file_handle.close()
        
        try:
           assert_file_equal(expect_out, real_out)        
        except AssertionError, e:
           print >>sys.stdout, "case %s failed" % self.casename
           print >>sys.stdout, "%s not equal to %s" % (real_out, expect_out)
           sys.exit(1)
        print >>sys.stdout, "case %s passed" % self.casename

if __name__ == "__main__":
     test = Test_rg_getsite_map('Test_rg_getsite_map', sample, 'rg_getsite', '/user/spider/rankget/input/', '/user/spider/rankget/output/site_to_train/')
     setup()
     test.set_input()
     test.run()
     test.check_result()  
     teardown()

     test1 = Test_rg_getsite('Test_rg_getsite', sample, 'rg_getsite', '/user/spider/rankget/input/', '/user/spider/rankget/output/site_to_train/')
     test1.set_input()
     test1.run()
     test1.check_result()  
